"""
提示词模板库
包含场景生成和Manim脚本生成的提示词模板，以及few-shot示例
"""

from typing import List, Dict, Any


class SceneGenerationTemplate:
    """场景生成提示词模板"""

    @staticmethod
    def get_system_prompt() -> str:
        return (
            "你是一名专业的教学视频脚本结构化助手。你的任务是将Markdown文档转换为结构化的视频场景数据。"
            "请严格遵循给定的JSON格式，只输出JSON，不要任何解释或额外文本。"
        )

    @staticmethod
    def get_user_prompt(markdown_text: str, scene_level: int = 2) -> str:
        return (
            f"请将下面的Markdown文档转换为用于生成视频的场景JSON。\n"
            f"使用{scene_level}级标题作为场景分隔符。\n\n"
            "严格遵循以下结构：\n"
            "{\n"
            '  "version": "1.0",\n'
            '  "meta": {"title": "视频标题", "language": "zh-CN"},\n'
            "  scenes: [\n"
            "    {\n"
            '      "id": "scene-001",\n'
            '      "title": "场景标题",\n'
            '      "narration": "分镜讲述文本，口语化表达",\n'
            '      "duration_hint_sec": 8.0,\n'
            '      "visuals": [\n'
            '        {"type": "text", "content": "要显示的文本", "position": "center", "effects": ["write"]}\n'
            "      ]\n"
            "    }\n"
            "  ]\n"
            "}\n\n"
            "输出要求：\n"
            "1. 只输出JSON（不需要解释、不要代码块围栏）\n"
            "2. narration应口语化，适合语音朗读\n"
            "3. 每个场景时长建议5-15秒\n"
            "4. visuals描述要显示的视觉元素\n\n"
            f"Markdown文档：\n{markdown_text}\n"
        )

    @staticmethod
    def get_few_shot_examples() -> List[Dict[str, str]]:
        """获取few-shot示例"""
        return [
            {
                "input": "# Python基础教程\n\n## 变量与数据类型\nPython是一种解释型语言...\n\n## 控制流程\n条件语句和循环是...\n",
                "output": """{
  "version": "1.0",
  "meta": {
    "title": "Python基础教程",
    "language": "zh-CN"
  },
  "scenes": [
    {
      "id": "scene-001",
      "title": "课程介绍",
      "narration": "欢迎来到Python基础教程，今天我们来学习Python的基本概念。",
      "duration_hint_sec": 8.0,
      "visuals": [
        {"type": "text", "content": "Python基础教程", "position": "center", "effects": ["write"]}
      ]
    },
    {
      "id": "scene-002",
      "title": "变量与数据类型",
      "narration": "首先，我们来了解Python中的变量和数据类型。变量就像是存储数据的容器。",
      "duration_hint_sec": 10.0,
      "visuals": [
        {"type": "text", "content": "变量 = 数据的容器", "position": "center", "effects": ["fade_in"]}
      ]
    }
  ]
}"""
            },
            {
                "input": "# 机器学习入门\n\n## 什么是机器学习\n机器学习是人工智能的一个分支...\n\n## 监督学习\n监督学习使用标记数据...\n",
                "output": """{
  "version": "1.0",
  "meta": {
    "title": "机器学习入门",
    "language": "zh-CN"
  },
  "scenes": [
    {
      "id": "scene-001",
      "title": "开场",
      "narration": "今天我们来探索机器学习的世界，看看AI是如何从数据中学习的。",
      "duration_hint_sec": 9.0,
      "visuals": [
        {"type": "text", "content": "机器学习入门", "position": "center", "effects": ["write"]}
      ]
    },
    {
      "id": "scene-002",
      "title": "机器学习概念",
      "narration": "机器学习简单来说，就是让计算机通过分析大量数据，自动发现规律和模式。",
      "duration_hint_sec": 12.0,
      "visuals": [
        {"type": "text", "content": "数据 → 算法 → 规律", "position": "center", "effects": ["create"]}
      ]
    }
  ]
}"""
            }
        ]


class ManimGenerationTemplate:
    """Manim脚本生成提示词模板"""

    @staticmethod
    def get_system_prompt() -> str:
        return (
            "你是一名专业的Manim动画工程师。你的任务是根据场景信息和字幕，生成高质量的Manim动画脚本。"
            "请只输出Python代码，不要任何解释、注释或额外文本。"
        )

    @staticmethod
    def get_user_prompt(
        scenes: List[Dict[str, Any]],
        captions: str,
        resolution: str,
        fps: int
    ) -> str:
        scene_descriptions = []
        for s in scenes:
            title = (s.get("title") or "").strip()
            narration = (s.get("narration") or "").strip()
            visuals = s.get("visuals", [])
            visual_desc = ", ".join([v.get("content", "") for v in visuals if v.get("content")])

            desc = f"- 标题: {title}"
            if narration:
                desc += f" | 讲述: {narration}"
            if visual_desc:
                desc += f" | 视觉: {visual_desc}"
            scene_descriptions.append(desc)

        return (
            "请根据以下场景信息与字幕，生成一个完整的Manim动画脚本。\n\n"
            "要求：\n"
            "1. 必须输出完整的Python代码，包含 `from manim import *` 和 `class Generated(Scene)`\n"
            "2. 在 `construct` 方法中实现动画逻辑\n"
            "3. 禁止任何网络/文件操作，不要读取外部资源\n"
            "4. 使用 Text、MathTex、Circle、Rectangle 等基础元素\n"
            "5. 场景间要有适当的过渡效果（FadeOut/FadeIn）\n"
            "6. 字体大小适中，确保在目标分辨率下清晰可读\n"
            f"7. 目标视频参数：分辨率 {resolution}，帧率 {fps}fps\n\n"
            "场景信息：\n" + "\n".join(scene_descriptions) + "\n\n"
            f"字幕(SRT格式)：\n{captions}\n"
        )

    @staticmethod
    def get_few_shot_examples() -> List[Dict[str, str]]:
        """获取few-shot示例"""
        return [
            {
                "input": """场景信息：
- 标题: Python变量 | 讲述: 变量就像是存储数据的容器
- 标题: 数据类型 | 讲述: Python有多种数据类型，包括数字、字符串、列表等

字幕：
1
00:00:00,000 --> 00:00:03,000
变量就像是存储数据的容器

2
00:00:03,500 --> 00:00:06,500
Python有多种数据类型""",
                "output": """from manim import *

class Generated(Scene):
    def construct(self):
        # 场景1: Python变量
        title1 = Text("Python变量", font_size=48, color=BLUE)
        desc1 = Text("变量就像是存储数据的容器", font_size=32)
        desc1.next_to(title1, DOWN)

        self.play(Write(title1))
        self.play(Write(desc1))
        self.wait(1.5)
        self.play(FadeOut(title1), FadeOut(desc1))

        # 场景2: 数据类型
        title2 = Text("数据类型", font_size=48, color=GREEN)
        desc2 = Text("数字、字符串、列表...", font_size=32)
        desc2.next_to(title2, DOWN)

        self.play(Write(title2))
        self.play(Write(desc2))
        self.wait(1.5)
"""
            },
            {
                "input": """场景信息：
- 标题: 机器学习概述 | 讲述: 机器学习是让计算机从数据中自动学习
- 标题: 监督学习 | 讲述: 使用标记数据训练模型的过程

字幕：
1
00:00:00,000 --> 00:00:03,000
机器学习是让计算机从数据中自动学习

2
00:00:03,500 --> 00:00:07,000
监督学习使用标记数据训练模型""",
                "output": """from manim import *

class Generated(Scene):
    def construct(self):
        # 场景1: 机器学习概述
        title1 = Text("机器学习", font_size=48, color=BLUE)
        arrow = Arrow(LEFT, RIGHT, color=YELLOW)
        desc1 = Text("数据 → 规律", font_size=36)
        group1 = VGroup(title1, arrow, desc1).arrange(RIGHT)

        self.play(Write(title1))
        self.play(GrowArrow(arrow))
        self.play(Write(desc1))
        self.wait(2)
        self.play(FadeOut(group1))

        # 场景2: 监督学习
        title2 = Text("监督学习", font_size=48, color=GREEN)
        input_box = Rectangle(width=2, height=1, color=BLUE)
        output_box = Rectangle(width=2, height=1, color=RED)
        arrow2 = Arrow(input_box.get_right(), output_box.get_left(), color=YELLOW)
        input_label = Text("标记数据", font_size=24).move_to(input_box.get_center())
        output_label = Text("模型", font_size=24).move_to(output_box.get_center())

        group2 = VGroup(input_box, output_label, arrow2, output_box, input_label)

        self.play(Write(title2))
        self.play(Create(input_box), Write(input_label))
        self.play(Create(arrow2))
        self.play(Create(output_box), Write(output_label))
        self.wait(2)
"""
            }
        ]


class PromptManager:
    """提示词管理器"""

    def __init__(self):
        self.scene_template = SceneGenerationTemplate()
        self.manim_template = ManimGenerationTemplate()

    def get_scene_generation_prompt(
        self,
        markdown_text: str,
        scene_level: int = 2,
        use_few_shot: bool = True
    ) -> List[Dict[str, str]]:
        """
        获取场景生成提示

        Args:
            markdown_text: Markdown文本
            scene_level: 场景标题层级
            use_few_shot: 是否使用few-shot示例

        Returns:
            格式为 [{"role": "system", "content": ...}, {"role": "user", "content": ...}] 的消息列表
        """
        messages = [
            {"role": "system", "content": self.scene_template.get_system_prompt()},
        ]

        if use_few_shot:
            examples = self.scene_template.get_few_shot_examples()
            for example in examples:
                messages.append({"role": "user", "content": f"示例输入:\n{example['input']}"})
                messages.append({"role": "assistant", "content": example['output']})

        messages.append({
            "role": "user",
            "content": self.scene_template.get_user_prompt(markdown_text, scene_level)
        })

        return messages

    def get_manim_generation_prompt(
        self,
        scenes: List[Dict[str, Any]],
        captions: str,
        resolution: str,
        fps: int,
        use_few_shot: bool = True
    ) -> List[Dict[str, str]]:
        """
        获取Manim脚本生成提示

        Args:
            scenes: 场景列表
            captions: 字幕文本
            resolution: 分辨率
            fps: 帧率
            use_few_shot: 是否使用few-shot示例

        Returns:
            格式为 [{"role": "system", "content": ...}, {"role": "user", "content": ...}] 的消息列表
        """
        messages = [
            {"role": "system", "content": self.manim_template.get_system_prompt()},
        ]

        if use_few_shot:
            examples = self.manim_template.get_few_shot_examples()
            for example in examples:
                messages.append({"role": "user", "content": f"示例输入:\n{example['input']}"})
                messages.append({"role": "assistant", "content": example['output']})

        messages.append({
            "role": "user",
            "content": self.manim_template.get_user_prompt(scenes, captions, resolution, fps)
        })

        return messages
